/*
 * Handle kernel unexport symbols
 */
#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
#ifdef TK5
#include <linux/kallsyms.h>
#endif
#include <linux/ftrace.h>
#include "syms.h"

/*
 * Unexport variables
 */
/*
 * Unexport symbols
 */
struct rq __percpu *p_runqueues;
#ifdef TK5
sys_call_ptr_t *p_sys_call_table[NR_syscalls+1];
struct mutex *p_module_mutex;
#else
#ifdef CONFIG_X86_64
sys_call_ptr_t *p_sys_call_table[__NR_syscall_max+1];
#elif defined CONFIG_ARM64
syscall_fn_t *p_sys_call_table[__NR_syscalls];
#endif
#endif
/*
 * Scheduler hooks
 */
u64 (*p_stat_arch_irq_stat_cpu)(unsigned int cpu);
u64 (*p_stat_arch_irq_stat)(void);
void (*p_stat_do_syscall_64)(unsigned long nr, struct pt_regs *regs);
int (*p_idle_cpu)(int cpu);
void (*p_stat_sched_rqm_switch)(struct rq *rq, struct task_struct *prev, struct task_struct *next);
void (*p_stat_psi_task_switch)(struct task_struct *prev, struct task_struct *next, bool sleep);
struct rq *(*p_stat_finish_task_switch)(struct task_struct *prev);
struct task_struct * (*p_stat_pick_next_task)(struct rq *rq, struct task_struct *prev, struct rq_flags_stat *rf);
struct page *
(*p__alloc_pages_nodemask)(gfp_t gfp_mask, unsigned int order, int preferred_nid,
                                                         nodemask_t *nodemask);
blk_qc_t (*p_submit_bio)(struct bio *bio);
void (*p_bio_endio)(struct bio *bio);
void *(*p__kmalloc)(size_t size, gfp_t flags);
void *(*p__kmalloc_node)(size_t size, gfp_t flags, int node);
void *(*p_kmem_cache_alloc)(struct kmem_cache *s, gfp_t gfpflags);
#ifndef TK5
bool (*p_blk_mq_get_driver_tag)(struct request *rq);
#else
bool (*p_blk_mq_dispatch_rq_list)(struct blk_mq_hw_ctx *hctx, struct list_head *list,
       unsigned int nr_budgets);
#endif
void (*p_do_page_fault)(struct pt_regs *regs, unsigned long error_code, unsigned long address);
void (*p_stat_rcu_note_context_switch)(bool preempt);
int (*p_stat_access_remote_vm)(struct mm_struct *mm, unsigned long addr,
         void *buf, int len, unsigned int gup_flags);
#ifdef TK5
vm_fault_t (*p_handle_mm_fault)(struct vm_area_struct *vma, unsigned long address,
                unsigned int flags, struct pt_regs *regs);
#endif

struct ksym
{
	int type;
	char *name;
	void **address;
};

enum ksym_type {
	KSYM_DEF,
	KSYM_JMP_MCOUNT
};
static struct ksym syms[] = {
#ifdef TK5
	{KSYM_DEF, "module_mutex", (void **)&p_module_mutex},
	{KSYM_JMP_MCOUNT, "__alloc_pages", (void **)&p__alloc_pages_nodemask},
	{KSYM_JMP_MCOUNT, "psi_task_switch", (void **)&p_stat_psi_task_switch},
	{KSYM_JMP_MCOUNT, "handle_mm_fault", (void **)&p_handle_mm_fault},
	{KSYM_JMP_MCOUNT, "blk_mq_dispatch_rq_list", (void **)&p_blk_mq_dispatch_rq_list},
#else
#ifdef CONFIG_X86_64
	{KSYM_JMP_MCOUNT, "do_syscall_64", (void **)&p_stat_do_syscall_64},
#endif
	{KSYM_JMP_MCOUNT, "__alloc_pages_nodemask", (void **)&p__alloc_pages_nodemask},
	{KSYM_JMP_MCOUNT, "finish_task_switch", (void **)&p_stat_finish_task_switch},
	{KSYM_JMP_MCOUNT, "blk_mq_get_driver_tag", (void **)&p_blk_mq_get_driver_tag},
#ifndef TK3
	{KSYM_JMP_MCOUNT, "do_page_fault", (void **)&p_do_page_fault},
#endif
#ifdef TK4_NEW
	{KSYM_JMP_MCOUNT, "sched_rqm_switch", (void **)&p_stat_sched_rqm_switch},
#else
	{KSYM_JMP_MCOUNT, "rcu_note_context_switch", (void **)&p_stat_rcu_note_context_switch},
#endif
#endif
	{KSYM_DEF, "runqueues", (void **)&p_runqueues},
	{KSYM_DEF, "sys_call_table", (void **)&p_sys_call_table},
	{KSYM_JMP_MCOUNT, "idle_cpu", (void **)&p_idle_cpu},
#if defined(TK4_OLD) || defined(TK3) || CONFIG_ARM64
	{KSYM_JMP_MCOUNT, "pick_next_task_fair", (void **)&p_stat_pick_next_task},
#else
	{KSYM_JMP_MCOUNT, "pick_next_task", (void **)&p_stat_pick_next_task},
#endif
	{KSYM_JMP_MCOUNT, "submit_bio", (void **)&p_submit_bio},
	{KSYM_JMP_MCOUNT, "bio_endio", (void **)&p_bio_endio},
#ifdef CONFIG_X86_64
	{KSYM_JMP_MCOUNT, "arch_irq_stat_cpu", (void **)&p_stat_arch_irq_stat_cpu},
	{KSYM_JMP_MCOUNT, "arch_irq_stat", (void **)&p_stat_arch_irq_stat},
#endif
	{KSYM_JMP_MCOUNT, "access_remote_vm", (void **)&p_stat_access_remote_vm},
	{KSYM_JMP_MCOUNT, "__kmalloc", (void **)&p__kmalloc},
	{KSYM_JMP_MCOUNT, "__kmalloc_node", (void **)&p__kmalloc_node},
	{KSYM_JMP_MCOUNT, "kmem_cache_alloc", (void **)&p_kmem_cache_alloc},
	{0, 0, 0}
};

int ksyms_init(void)
{
	unsigned long addr;
	int i = 0;

	/* Init kernel symbols */
	while (true) {
		if (!syms[i].name)
			break;

#ifdef TK5
		addr = kallsyms_lookup_name_tk5(syms[i].name);
#else
		addr = kallsyms_lookup_name(syms[i].name);
#endif
		if (!addr) {
			pr_err("symbol %s not found\n", syms[i].name);
			return -ENODEV;
		}

		switch (syms[i].type) {
		case KSYM_DEF:
			syms[i++].address[0] = (void *)addr;
			break;
		case KSYM_JMP_MCOUNT:
			syms[i++].address[0] = (void *)(addr + MCOUNT_INSN_SIZE);
			break;
		default:
			pr_err("symbol %s invalid type %d\n", syms[i].name, syms[i].type);
			return -EINVAL;
		}
	}
	return 0;
}
